package com.gitee.yanfanvip.cluster;

import org.jgroups.Address;
import org.jgroups.ChannelListener;
import org.jgroups.Event;
import org.jgroups.Global;
import org.jgroups.Header;
import org.jgroups.JChannel;
import org.jgroups.Message;
import org.jgroups.View;
import org.jgroups.annotations.MBean;
import org.jgroups.conf.ClassConfigurator;
import org.jgroups.stack.Protocol;
import org.jgroups.util.MessageBatch;
import org.rocksdb.RocksDBException;
import com.gitee.yanfanvip.interfaces.Database;
import com.gitee.yanfanvip.model.OptionsEnum;
import com.gitee.yanfanvip.model.Row;
import com.gitee.yanfanvip.single.SingleDatabase;
import java.io.BufferedOutputStream;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Supplier;

@MBean(description="集群化数据库")
public class ClusterDatabase extends Protocol implements Database, ClusterReceiver{
	protected static final short ClusterDatabase_ID = 11501;
	protected static final short ClusterDatabase_Header_ID = 21501;
	static{
		ClassConfigurator.add(ClusterDatabase_Header_ID, ClusterDatabaseHeader.class);
		ClassConfigurator.addProtocol(ClusterDatabase_ID, ClusterDatabase.class);
	}
	SingleDatabase database;
	String path = null;
	protected Address localAddress;
	protected List<Address> members = new ArrayList<>();
	protected RemoteCluster remoteCluster;
	protected boolean enable = false;
	
	public ClusterDatabase(String cluster, int port) {
		path = cluster + "-" + port;
	}
	
	public void start() throws Exception {
		try {
			remoteCluster = getProtocolStack().findProtocol(RemoteCluster.class);
			database = SingleDatabase.get(path);
			JChannel channel = getProtocolStack().getChannel();
			channel.setReceiver(this);
			channel.addChannelListener(new ChannelListener() {
				@Override
				public void channelDisconnected(JChannel channel) { }
				
				@Override
				public void channelConnected(JChannel channel) {
					try {
						channel.getState(null, 60 * 60 * 1000, false);
					} catch (Exception e) {
						e.printStackTrace();
						System.exit(-1);
					}
				}
				
				@Override
				public void channelClosed(JChannel channel) { }
			});
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}
	
	@Override
    public void getState(OutputStream out) throws Exception {
        try(ObjectOutputStream output = new ObjectOutputStream(new BufferedOutputStream(out, 1024))) {
        	Set<String> keyspaces = database.keyspaces();
        	for (String keyspace : keyspaces) {
        		database.forEach(keyspace, (key, value)->{
        			try {
        				output.writeObject(new Row(keyspace, key, value));
					} catch (IOException e) {
						throw new RuntimeException(e);
					}
        		});
			}
        	output.writeObject(null);
        }
    }

    @Override
	public void setState(InputStream in) throws Exception {
    	database.clearAll();
        try(ObjectInputStream input = new ObjectInputStream(in)) {
        	while(true){
        		Row row = (Row) input.readObject();
        		if(row == null){ break; }
        		this.database.put(row.getKeyspaces(), row.getKey(), row.getValue());
        	}
        }
        this.enable = true;
    }
	
	@Override
	public Object up(Message msg) {
		ClusterDatabaseHeader header = msg.getHeader(id);
		if(header != null){
			if(msg.getSrc().equals(this.localAddress)){
				return super.up(msg);
			}
			switch (header.options) {
			case PUT:{
				Row row = msg.getObject();
				_put(row.getKeyspaces(), row.getKey(), row.getValue());
				break;
			}
			case PUTALL:{
				Set<Row> rows =  msg.getObject();
				for (Row row : rows) {
					_put(row.getKeyspaces(), row.getKey(), row.getValue());
				}
				break;
			}
			case REMOVE:{
				Row row = msg.getObject();
				_remove(row.getKeyspaces(), row.getKey());
				break;
			}
			case CLEAR:{
				String keyspaces = msg.getObject();
				_clear(keyspaces);
				break;
			}
			default:
				throw new RuntimeException("bad request");
			}
		}
		return super.up(msg);
	}
	
	@Override
	public void up(MessageBatch batch) {
		for (Message msg : batch) {
			ClusterDatabaseHeader header = msg.getHeader(id);
			if(header != null){
				batch.remove(msg);
				up(msg);
			}
		}
		if(!batch.isEmpty()){
			up_prot.up(batch);
		}
	}
	
	@Override
	public Object up(Event evt) {
		switch (evt.type()) {
			case Event.VIEW_CHANGE:{
				View view = evt.getArg();
				handleView(view);
				break;
			}
			case Event.SET_LOCAL_ADDRESS:
                localAddress=evt.getArg();
                break;
			case 11501:
				break;
		}
		return super.up(evt);
	}
	
	@Override
	public Object down(Event evt) {
		switch (evt.type()) {
			case Event.VIEW_CHANGE:{
				View view = evt.getArg();
				handleView(view);
				break;
			}
			case Event.SET_LOCAL_ADDRESS:
                localAddress=evt.getArg();
                break;
			case 11501:
				break;
		}
		return super.down(evt);
	}
	
	protected void handleView(View view) {
		List<Address> mbrs = view.getMembers();
		if(mbrs.size() == 1){ this.enable = true; }
		List<Address> added = new ArrayList<>();
		mbrs.forEach(m->{
			if(!members.contains(m)) {
				added.add(m);
			}
		});
		members = mbrs;
		handleAddView(added);
	}
	/**
	 * 新增的节点
	 * @param added 新增的节点地址
	 */
	protected void handleAddView(List<Address> added) { }

	@Override
	public boolean containsKey(String keyspaces, String key) throws RocksDBException {
		return database.containsKey(keyspaces, key);
	}

	@Override
	public byte[] get(String keyspaces, String key) throws RocksDBException {
		return database.get(keyspaces, key);
	}

	@Override
	public void forEach(String keyspaces, BiConsumer<String, byte[]> consumer) throws RocksDBException {
		database.forEach(keyspaces, consumer);
	}
	
	public Set<String> keyspaces() throws RocksDBException {
		return this.database.keyspaces();
	}
	
	@Override
	public void put(String keyspaces, String key, byte[] value) throws RocksDBException {
		_cluster_put(keyspaces, key, value);
		remoteCluster.put(keyspaces, key, value);
	}
	
	public void _cluster_put(String keyspaces, String key, byte[] value) throws RocksDBException {
		sendMessage(OptionsEnum.PUT, new Row(keyspaces, key, value));
		_put(keyspaces, key, value);
	}

	@Override
	public void putAll(String keyspaces, Map<String, byte[]> datas) throws RocksDBException {
		_cluster_putAll(keyspaces, datas);
		remoteCluster.putAll(keyspaces, datas);
	}
	
	public void _cluster_putAll(String keyspaces, Map<String, byte[]> datas) throws RocksDBException {
		Set<Row> rows = new HashSet<Row>();
		for (String key : datas.keySet()) {
			byte[] value = datas.get(key);
			rows.add(new Row(keyspaces, key, value));
		}
		sendMessage(OptionsEnum.PUTALL, rows);
		_putAll(keyspaces, datas);
	}

	@Override
	public void remove(String keyspaces, String key) throws RocksDBException {
		_cluster_remove(keyspaces, key);
		remoteCluster.remove(keyspaces, key);
	}

	public void _cluster_remove(String keyspaces, String key) throws RocksDBException {
		sendMessage(OptionsEnum.REMOVE, new Row(keyspaces, key));
		_remove(keyspaces, key);
	}
	
	@Override
	public void clear(String keyspaces) throws RocksDBException {
		_cluster_clear(keyspaces);
		remoteCluster.clear(keyspaces);
	}
	
	public void _cluster_clear(String keyspaces) throws RocksDBException {
		sendMessage(OptionsEnum.CLEAR, keyspaces);
		_clear(keyspaces);
	}
	
	public void _put(String keyspaces, String key, byte[] value) {
		try {
			database.put(keyspaces, key, value);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _putAll(String keyspaces, Map<String, byte[]> datas) {
		try {
			database.putAll(keyspaces, datas);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _remove(String keyspaces, String key) {
		try {
			database.remove(keyspaces, key);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}

	public void _clear(String keyspaces) {
		try {
			database.clear(keyspaces);
		} catch (RocksDBException e) {
			throw new RuntimeException(e);
		}
	}
	
	void sendMessage(OptionsEnum option, Object data){
		Message msg = new Message(null, data);
		msg.putHeader(id, new ClusterDatabaseHeader(option));
		down_prot.down(msg);
	}

	@Override
	public void close() throws Exception {
		database.close();
	}
	
	public SingleDatabase getSingleDatabase() {
		return database;
	}
	
	public static class ClusterDatabaseHeader extends Header{
	    public Supplier<? extends Header> create() { return ClusterDatabaseHeader::new; }
	    
	    private OptionsEnum options;
	    
	    public ClusterDatabaseHeader() {
	    	super();
	    }
	    
	    public ClusterDatabaseHeader(OptionsEnum options) {
	    	this.options = options;
	 	}
	    
	    public short getMagicId() {return ClusterDatabase_Header_ID;}
	    
	    public int serializedSize() { return Global.INT_SIZE; }
	    
	    public String toString() {
	        return options.name();
	    }

	    public void writeTo(DataOutput out) throws Exception {
	    	out.writeInt(options.getType());
	    }

	    public void readFrom(DataInput in) throws Exception {
	    	this.options = OptionsEnum.get(in.readInt());
	    }
	}

	public void forEach(String keyspaces, int batch, Consumer<Set<Row>> consumer) throws RocksDBException {
		Set<Row> batchResult = new HashSet<Row>(batch);
		database.forEach(keyspaces, (key, value)->{
			batchResult.add(new Row(keyspaces, key, value));
			if(batchResult.size() > 10000){
				consumer.accept(batchResult);
				batchResult.clear();
			}
		});
		if(batchResult.size() > 0){
			consumer.accept(batchResult);
			batchResult.clear();
		}
	}

}
